from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain.prompts import ChatPromptTemplate
from langchain_community.llms import Ollama
from langchain_community.vectorstores import Chroma
from langchain.schema import Document
from typing import Dict, Any
import os


class DocumentGenerationChain:
    """文档生成处理链"""

    def __init__(self, llm: Ollama, vector_store: Chroma):
        """初始化文档生成链

        Args:
            llm: 用于生成文档的语言模型
            vector_store: 用于检索相关上下文的向量存储
        """
        self.llm = llm
        self.vector_store = vector_store
        self.output_parser = StrOutputParser()
        self.chain = None

        # 加载提示模板
        template_path = os.path.join(os.path.dirname(__file__),
                                     "prompt_templates", "doc_template.txt")
        with open(template_path, "r", encoding="utf-8") as f:
            self.template = f.read()

    def setup_chain(self):
        """设置处理链"""
        # 文档生成提示模板
        documentation_prompt = ChatPromptTemplate.from_template(self.template)

        # 创建检索器
        retriever = self.vector_store.as_retriever()

        # 完整处理链
        self.chain = (
                {"code": RunnablePassthrough(), "context": retriever}
                | documentation_prompt
                | self.llm
                | self.output_parser
        )

    def process_file(self, file_content: str) -> str:
        """处理单个文件并生成文档

        Args:
            file_content: 文件内容

        Returns:
            生成的文档内容

        Raises:
            ValueError: 如果处理链未设置
        """
        if not self.chain:
            raise ValueError("请先调用setup_chain()设置处理链")

        return self.chain.invoke(file_content)

    def process_project(self, file_contents: Dict[str, str]) -> Dict[str, str]:
        """处理整个项目的文件并生成文档

        Args:
            file_contents: 文件路径到文件内容的映射

        Returns:
            文件路径到生成文档的映射
        """
        if not self.chain:
            raise ValueError("请先调用setup_chain()设置处理链")

        results = {}
        for file_path, content in file_contents.items():
            print(f"正在处理文件: {file_path}")
            results[file_path] = self.process_file(content)

        return results
